from pathlib import Path
import cv2
import numpy as np
from torch.utils.data import Dataset


CUR_DIR = data_path = Path(__file__).resolve().parent


class CUB(Dataset):
    def __init__(self, data_path=CUR_DIR / '../dataset/CUB_200_2011/images', train=True, train_rate=0.8, max_classes=0):
        super().__init__()
        self.classes = []
        x, y = [], []
        for label_dir in Path(data_path).iterdir():
            if not label_dir.is_dir():
                continue
            self.classes.append(label_dir.name)
            for img_path in label_dir.iterdir():
                print('Loaded %d images' % len(y), end='\r')
                img = cv2.imread(str(img_path))
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                img = cv2.resize(img, (224, 224), interpolation=cv2.INTER_LINEAR)
                img = img.transpose(2, 0, 1)
                x.append(img)
                y.append(len(self.classes) - 1)
            if len(self.classes) == max_classes:  # Use only the top `max_classes` categories if `max_classes > 0`.
                break
        x = np.array(x, dtype=np.float32)
        x = cv2.normalize(x, None, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX)
        y = np.array(y, dtype=np.int64)

        num_train = int(len(y) * train_rate)
        np.random.seed(197)
        rand_indices = np.random.permutation(len(y))
        if train:
            self.x = x[rand_indices][:num_train]
            self.y = y[rand_indices][:num_train]
        else:
            self.x = x[rand_indices][num_train:]
            self.y = y[rand_indices][num_train:]

    def __getitem__(self, index):
        return self.x[index], self.y[index]  # An image (3,244,244) and its label

    def __len__(self):
        return len(self.y)
